package com.boat.framework.common.thread.pool;

import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;

import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import com.boat.framework.common.thread.context.AppContext;
import com.boat.framework.common.thread.context.ContextUtils;
import com.boat.framework.common.thread.context.LogContext;
import com.boat.framework.common.thread.context.MockContext;

/**
 * 用于父子线程复制上下文信息
 * 
 * @author 周江
 *
 */
public class ContextThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {
	private static final long serialVersionUID = 6147231718113027677L;

	/**
	 * 把父线程的上下文信息赋值给子线程
	 * 
	 * @param runnable
	 */
	@Override
	public void execute(final Runnable runnable) {
		// 获取父线程上下文
		final Map<String, Object> appContextMap = AppContext.getCopyOfContextMap();
		final Map<String, String> logContextMap = LogContext.getCopyOfContextMap();
		final Map<String, Object> mockContextMap = MockContext.getCopyOfContextMap();

		super.execute(new Runnable() {
			@Override
			public void run() {
				// 将父线程上下文复制给子线程
				if (appContextMap != null) {
					AppContext.setContextMap(appContextMap);
				}
				if (logContextMap != null) {
					LogContext.setContextMap(logContextMap);
				}
				if (mockContextMap != null) {
					MockContext.setContextMap(mockContextMap);
				}

				try {
					runnable.run();
				} finally {
					// 清空线程上下文内容
					ContextUtils.clear();
				}
			}
		});
	}

	/**
	 * 把父线程的上下文信息赋值给子线程
	 * 
	 * @param Callable
	 */
	@Override
	public <T> Future<T> submit(final Callable<T> task) {
		// 获取父线程上下文
		final Map<String, Object> appContextMap = AppContext.getCopyOfContextMap();
		final Map<String, String> logContextMap = LogContext.getCopyOfContextMap();
		final Map<String, Object> mockContextMap = MockContext.getCopyOfContextMap();
		
		return super.submit(new Callable<T>() {
			@Override
			public T call() throws Exception {
				// 将父线程上下文复制给子线程
				if (appContextMap != null) {
					AppContext.setContextMap(appContextMap);
				}
				if (logContextMap != null) {
					LogContext.setContextMap(logContextMap);
				}
				if (mockContextMap != null) {
					MockContext.setContextMap(mockContextMap);
				}
				
				try {
					return task.call();
				} finally {
					// 清空线程上下文内容
					ContextUtils.clear();
				}
			}
		});
	}
	
}